import torch
from torchvision import transforms
from PIL import Image

__all__ = ['personTrainTransform', 'personTestTransform']

IMG_H = 256
IMG_W = 128
p = 0.5

personTrainTransform = transforms.Compose([
    transforms.Resize((IMG_H, IMG_W)),
    transforms.RandomHorizontalFlip(p),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.485, 0.456, 0.406),
                         std=(0.229, 0.224, 0.225)),
])

# personTestTransform = transforms.Compose([
#     transforms.Resize((IMG_H, IMG_W)),
#     transforms.ToTensor(),
#     transforms.Normalize(mean=(0.485, 0.456, 0.406),
#                          std=(0.229, 0.224, 0.225)),
# ])


class RectScale(object):
    def __init__(self, height, width, interpolation=Image.BILINEAR):
        self.height = height
        self.width = width
        self.interpolation = interpolation

    def __call__(self, img):
        w, h = img.size
        if h == self.height and w == self.width:
            return img
        return img.resize((self.width, self.height), self.interpolation)


personTestTransform = transforms.Compose([
    RectScale(IMG_H, IMG_W),
    transforms.ToTensor(),
])
